package com.hdu;

import com.hdu.TimerTaskList.TimerTaskEntry;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.DelayQueue;
import java.util.concurrent.atomic.AtomicInteger;

import static java.util.Objects.nonNull;
import static java.util.Optional.ofNullable;

public class TimingWheel {

    private final Long TICK_MS;
    private final Integer WHEEL_SIZE;
    private final AtomicInteger TASK_COUNTER;
    private final DelayQueue<TimerTaskList> QUEUE;

    private final Long INTERVAL;
    private final List<TimerTaskList> BUCKETS;
    private Long currentTime;

    private TimingWheel overflowWheel = null;

    public TimingWheel(Long TICK_MS,
                       Integer WHEEL_SIZE,
                       Long startMs,
                       AtomicInteger TASK_COUNTER,
                       DelayQueue<TimerTaskList> QUEUE) {

        this.TICK_MS = TICK_MS;
        this.WHEEL_SIZE = WHEEL_SIZE;
        this.TASK_COUNTER = TASK_COUNTER;
        this.QUEUE = QUEUE;
        INTERVAL = TICK_MS * WHEEL_SIZE;
        currentTime = startMs - (startMs % TICK_MS);

        BUCKETS = new ArrayList<>(WHEEL_SIZE);
        for (int i = 0; i < WHEEL_SIZE; i++) {
            BUCKETS.add(new TimerTaskList(TASK_COUNTER));
        }
    }

    public synchronized void createOverFlowWheel() {
        if (overflowWheel == null) {
            overflowWheel = new TimingWheel(
                    INTERVAL,
                    WHEEL_SIZE,
                    currentTime,
                    TASK_COUNTER,
                    QUEUE
            );
        }
    }

    public boolean addTask(TimerTaskEntry timerTaskEntry) {
        Long expiration = timerTaskEntry.expirationMs;

        if (timerTaskEntry.cancel()) {
            return false;
        } else if (expiration < currentTime + TICK_MS) {
            return false;
        } else if (expiration < currentTime + INTERVAL) {
            Long virtualId = expiration / TICK_MS;
            TimerTaskList bucket = BUCKETS.get((int) (virtualId % WHEEL_SIZE));
            bucket.add(timerTaskEntry);
            if (bucket.setExpiration(virtualId * TICK_MS)) {
                QUEUE.offer(bucket);
            }
            return true;
        } else {
            return ofNullable(overflowWheel)
                    .map(wheel -> wheel.addTask(timerTaskEntry))
                    .orElseGet(() -> {
                        createOverFlowWheel();
                        return overflowWheel.addTask(timerTaskEntry);
                    });
        }
    }

    public void advanceClock(Long timeMs) {
        if (timeMs >= currentTime + TICK_MS) {
            currentTime = timeMs - (timeMs % TICK_MS);
            if (nonNull(overflowWheel)) {
                overflowWheel.advanceClock(currentTime);
            }
        }
    }
}